#! /usr/bin/env python3
# -*- coding: utf-8 -*-

from ESharkDSLVisitor import ESharkDSLVisitor
from ESharkDSLParser import ESharkDSLParser
from elasticsearch_dsl.query import Term, Q
from enum import Enum


class BoolQueryOp(Enum):
    """
    ElasticSearch 布尔操作符
    """
    MUST = 1
    MUST_NOT = 2
    SHOULD = 3


class MyVisitor(ESharkDSLVisitor):
    bqo_desc = {
        BoolQueryOp.MUST: 'must',
        BoolQueryOp.MUST_NOT: 'must_not',
        BoolQueryOp.SHOULD: 'should'
    }

    ops = {
        ESharkDSLParser.AND: BoolQueryOp.MUST,
        ESharkDSLParser.NOT: BoolQueryOp.MUST_NOT,
        ESharkDSLParser.OR: BoolQueryOp.SHOULD,
    }

    def __init__(self):
        self.querys = {
            '': [],
            'bool': {
                'must': [],
                'must_not': [],
                'should': []
            }
        }
        self.is_not = False
        self.is_op = -1

    def get_is_not(self):
        """
        用完即焚，避免影响之后的not操作符
        :return:
        """
        result = self.is_not
        self.is_not = False
        return result

    def visitProtocol(self, ctx):
        is_not = ctx.NOT()
        value = ctx.PROTOCOL().getText()
        q = Term(layers=value)

        if is_not:
            q = ~q

        if self.get_is_not():
            q = ~q

            if self.is_op == ESharkDSLParser.AND:
                # must -> must_not
                self.is_op = ESharkDSLParser.NOT

        if self.is_op > 0:
            op_name = self.bqo_desc[self.ops[self.is_op]]
            self.querys['bool'][op_name].append(q)
        else:
            self.querys[''].append(q)

    def visitProtocolFieldNameOpValue(self, ctx):
        name = ctx.FIELD_NAME().getText()
        value = ctx.FIELD_VALUE().getText()
        op_eq = ctx.op.type == ESharkDSLParser.EQ
        op_ne = ctx.op.type == ESharkDSLParser.NE
        op_gt = ctx.op.type == ESharkDSLParser.GT
        op_lt = ctx.op.type == ESharkDSLParser.LT
        op_ge = ctx.op.type == ESharkDSLParser.GE
        op_le = ctx.op.type == ESharkDSLParser.LE

        q = Q()

        if op_eq:
            q = Q({"term": {name: value}})
        elif op_ne:
            q = ~Q({"term": {name: value}})
        elif op_gt:
            q = Q({"range": {name: {'gt': value}}})
        elif op_lt:
            q = Q({"range": {name: {'lt': value}}})
        elif op_ge:
            q = Q({"range": {name: {'gte': value}}})
        elif op_le:
            q = Q({"range": {name: {'lte': value}}})
        else:
            pass

        if self.is_op == ESharkDSLParser.AND and self.get_is_not():
            # must -> must_not
            self.is_op = ESharkDSLParser.NOT

        if self.get_is_not():
            q = ~q

        if self.is_op > 0:
            op_name = self.bqo_desc[self.ops[self.is_op]]
            self.querys['bool'][op_name].append(q)
        else:
            self.querys[''].append(q)

    def visitExprOpexpr(self, ctx):
        if ctx.AND():
            self.is_op = ESharkDSLParser.AND
        elif ctx.OR():
            self.is_op = ESharkDSLParser.OR

        self.visit(ctx.expr(0))
        self.visit(ctx.expr(1))

    def visitParens(self, ctx):
        self.is_not = bool(ctx.NOT())
        return self.visit(ctx.expr())
